'''
根据name 提供不同算法对象（layer1 （first is 1））
和业务逻辑无关
'''
import cv2
import numpy as np

from algorithom.Yolov5_insSeg.insSeg_yolov5 import Inference as SegV5Inference
from algorithom.bar_pyzbar.bar import get_bar  # 条码识别
from algorithom.yolov8.user_det_predict import get_v8_detect_onepic  # v8 det
from algorithom.yolov8.user_seg_predict import get_yolov8_seg_predict # v8 seg infer


def get_seg_yolov5(model_path='', input_h=416, input_w=416, conf=0.25,
                   iou=0.45):  #

    infer = SegV5Inference(model_address=model_path, input_h=input_h, input_w=input_w, conf=conf, iou=iou)

    def seg_yolov5(img, isRoi = False, target_clsid = 0 ):
        '''

        :param img:
        :return: pre: n*(x1,y1,x2,y2,conf,cls) numpy  masks: h*w*n numpy
        '''
        H,W = img.shape[:2]
        pre, masks = infer.infer(img)# 第一次推理
        # img_show = infer.draw_box_mask(img, pre, masks)  # 画图
        img_show = img.copy()
        # clsid = 0  # 类别
        target_inds = np.argwhere(pre[:, 5] == target_clsid)
        target_inds = target_inds.reshape(-1).tolist()
        if isRoi and len(target_inds)>0:

            x1, y1, x2, y2 = pre[target_inds[0], 0:4] # 取目标类别置信度最大的
            padding = max(x2-x1,y2-y1)/8
            box_clip = (int(max(x1-padding,0)), int(max(y1-padding,0)), int(min(x2+padding,W)), int(min(y2+padding,H)))
            x1_clip, y1_clip, x2_clip, y2_clip = box_clip
            cv2.rectangle(img_show, box_clip[:2], box_clip[2:], (255, 0, 0), 20)
            img_clip = img[y1_clip:y2_clip, x1_clip:x2_clip]
            pre_mini, masks_mini = infer.infer(img_clip) # 第二次推理
            # 缩放到原图
            # pre: n*(x1,y1,x2,y2,conf,cls) numpy  masks: h*w*n numpy
            pre_mini[:,0] = pre_mini[:,0] + box_clip[0] # x1
            pre_mini[:, 1] = pre_mini[:, 1] + box_clip[1] # y1
            pre_mini[:,2] = pre_mini[:,2] + box_clip[0] # x2
            pre_mini[:, 3] = pre_mini[:, 3] + box_clip[1] # y2

            masks_mini = np.pad(masks_mini, ((y1_clip, H-y2_clip), (x1_clip, W-x2_clip),(0,0)), 'constant', constant_values=(0, 0)) # 填充
            img_show = infer.draw_box_mask(img_show, pre_mini, masks_mini)  # 画mini结果到原图
            return img_show, pre_mini, masks_mini  # user
        img_show = infer.draw_box_mask(img_show, pre, masks)  # 画图
        return img_show, pre, masks

    return seg_yolov5  # 返回推理函数

class AlgoFactory: #
    @staticmethod
    def create(class_type, **kwargs):
        if class_type == "yolov5_seg":
            seg_yolov5 = get_seg_yolov5(model_path=kwargs.get('model_path',' '),
                                        input_h=kwargs.get('input_h', 416),
                                        input_w=kwargs.get('input_w', 416),
                                        conf=kwargs.get('conf', 0.25),
                                        iou=kwargs.get('iou', 0.45))
            return seg_yolov5
        elif class_type == "yolov8_det":
            return get_v8_detect_onepic(
                model_path=kwargs.get('model_path', ' '),
            )
        elif class_type == "yolov8_seg":
            return get_yolov8_seg_predict(
                model_path=kwargs.get('model_path', ' '),
            )

        elif class_type == "bar_pyzbar":
            return get_bar
        else:
            raise ValueError("Invalid transport type")